#include "image.h"
#include "../../base/obj_cache.h"
#include "../../base/logger.h"
#include <png.h>
#include <jpeglib.h>
#include <fstream>

OPEN_JLIB_NS


struct PngReadContext {
  const u8* data;
  uwl size;
};

struct PngWriteContext {
  u8* data;
  uwl size;
  uwl iter;
};

static void __PNGRead (png_structp meta, png_bytep data, png_size_t size) noexcept {
  auto ctx = (PngReadContext*)png_get_io_ptr(meta);
  if (ctx->size < size) {
    log_e("__PNGRead(): Not enough data");
    longjmp(png_jmpbuf(meta), 1);
  }
  memcpy(data, ctx->data, size);
  ctx->data += size;
  ctx->size -= size;
}

static void __PNGWrite (png_structp meta, png_bytep data, png_size_t size) noexcept {
  auto ctx = (PngWriteContext*)png_get_io_ptr(meta);
  if (!ctx->data) {
    ctx->size = meta->width * meta->height * meta->channels;
    ctx->data = (u8*)malloc(ctx->size);
  }
  auto free = ctx->size - ctx->iter;
  if (free < size) {
    ctx->size = (ctx->iter + size) << 1;
    ctx->data = (u8*)realloc(ctx->data, ctx->size);
  }
  if (!ctx->data)
    longjmp(png_jmpbuf(meta), 1);
  memcpy(ctx->data + ctx->iter, data, size);
  ctx->iter += size;
}

static void __PNGError (png_structp meta, png_const_charp message) noexcept {
  log_e("__PNGError(): %s", message);
  longjmp(png_jmpbuf(meta), 1);
}

static void __PNGWarn (png_structp meta, png_const_charp message) noexcept {
  log_w("__PNGWarn(): %s", message);
}

static Texture* __loadPNG (const u8* data, uwl size) noexcept {
  PngReadContext ctx = {data, size};
  png_bytep pixels = nullptr;
  png_bytepp rowPointer = nullptr;
  png_infop info = nullptr;
  png_structp meta = png_create_read_struct(PNG_LIBPNG_VER_STRING, nullptr, __PNGError, __PNGWarn);
  if (!meta)
    return nullptr;

  if (setjmp(png_jmpbuf(meta))) {
    png_destroy_read_struct(&meta, &info, nullptr);
    if (pixels)
      free(pixels);
    if (rowPointer)
      free(rowPointer);
    return nullptr;
  }

  info = png_create_info_struct(meta);
  if (!info)
    longjmp(png_jmpbuf(meta), 1);
  png_set_read_fn(meta, &ctx, __PNGRead);
  png_read_info(meta, info);

  auto w = png_get_image_width(meta, info);
  auto h = png_get_image_height(meta, info);
  auto bitDepth = png_get_bit_depth(meta, info);
  auto colorType = png_get_color_type(meta, info);
  if (colorType == PNG_COLOR_TYPE_PALETTE)
    png_set_palette_to_rgb(meta);
  if (colorType == PNG_COLOR_TYPE_GRAY && bitDepth < 8)
    png_set_expand_gray_1_2_4_to_8(meta);
  if (png_get_valid(meta, info, PNG_INFO_tRNS))
    png_set_tRNS_to_alpha(meta);
  if (bitDepth == 16)
    png_set_strip_16(meta);
  if (bitDepth < 8)
    png_set_packing(meta);
  png_read_update_info(meta, info);

  GLenum fmt;
  colorType = png_get_color_type(meta, info);
  switch (colorType) {
    case PNG_COLOR_TYPE_GRAY:
      fmt = GL_LUMINANCE;
      break;
    case PNG_COLOR_TYPE_GRAY_ALPHA:
      fmt = GL_LUMINANCE_ALPHA;
      break;
    case PNG_COLOR_TYPE_RGB:
      fmt = GL_RGB;
      break;
    case PNG_COLOR_TYPE_RGB_ALPHA:
      fmt = GL_RGBA;
      break;
    default:
      log_e("__loadPNG(): Unknown color type");
      longjmp(png_jmpbuf(meta), 1);
  }

  auto rowBytes = png_get_rowbytes(meta, info);
  auto bytes = rowBytes * h;
  pixels = (png_bytep)malloc(bytes);
  if (!pixels)
    longjmp(png_jmpbuf(meta), 1);
  rowPointer = (png_bytepp)malloc(sizeof(png_bytep) * h);
  if (!rowPointer)
    longjmp(png_jmpbuf(meta), 1);
  for (s32 i = 0, j = (s32)h - 1; i < h; ++i, --j)
    rowPointer[i] = pixels + j * rowBytes;

  png_read_image(meta, rowPointer);
  png_read_end(meta, info);
  png_destroy_read_struct(&meta, &info, nullptr);

  auto texture = ObjCache::create<Texture>(w, h, fmt, GL_UNSIGNED_BYTE, pixels);
  free(rowPointer);
  free(pixels);
  return texture;
}

static ByteBuffer<> __savePNG (Texture* texture, s32 quality) noexcept {
  PngWriteContext ctx = {nullptr, 0, 0};
  png_bytep pixels = nullptr;
  png_bytep *rowPointers = nullptr;
  png_infop info = nullptr;
  png_structp meta = png_create_write_struct(PNG_LIBPNG_VER_STRING, nullptr, __PNGError, __PNGWarn);
  if (!meta)
    return {};

  if (setjmp(png_jmpbuf(meta))) {
    png_destroy_write_struct(&meta, &info);
    if (rowPointers)
      free(rowPointers);
    if (pixels)
      free(pixels);
    return {};
  }

  auto w = (u32)texture->width();
  auto rowBytes = 4UL * w;
  auto h = (u32)texture->height();
  auto bytes = rowBytes * h;
  pixels = (png_bytep)malloc(bytes);
  if (!pixels)
    longjmp(png_jmpbuf(meta), 1);
  glBindTexture(GL_TEXTURE_2D, texture->id());
  glGetTexImage(GL_TEXTURE_2D, 0, GL_RGBA, GL_UNSIGNED_BYTE, pixels);

  info = png_create_info_struct(meta);
  if (!info)
    longjmp(png_jmpbuf(meta), 1);

  rowPointers = (png_bytepp)malloc(sizeof(png_bytep) * h);
  if (!rowPointers)
    longjmp(png_jmpbuf(meta), 1);
  for (int i = 0, j = h - 1; i < h; ++i, --j)
    rowPointers[i] = (png_bytep)pixels + j * rowBytes;

  png_set_IHDR(meta, info, w, h, 8, PNG_COLOR_TYPE_RGBA, PNG_INTERLACE_NONE, PNG_COMPRESSION_TYPE_DEFAULT, PNG_FILTER_TYPE_DEFAULT);
  png_set_rows(meta, info, rowPointers);
  png_set_write_fn(meta, &ctx, __PNGWrite, nullptr);
  png_write_png(meta, info, PNG_TRANSFORM_IDENTITY, nullptr);
  png_write_end(meta, info);

  free(rowPointers);
  free(pixels);
  return {ctx.data, ctx.size, ctx.iter};
}



struct JpgReadContext {
  jmp_buf jump;
};

struct JpgWriteContext {
  jmp_buf jump;
  u8* data;
  uwl size;
  uwl iter;
};

static void __JPGError (j_common_ptr info) noexcept {
  auto ctx = (JpgReadContext*)info->client_data;
  char buffer[JMSG_LENGTH_MAX];
  info->err->format_message(info, buffer);
  log_e("__JPGError(): %s", buffer);
  longjmp(ctx->jump, 1);
}

static void __JPGWriteStart (j_compress_ptr info) noexcept {
  auto ctx = (JpgWriteContext*)info->client_data;
  ctx->size = info->image_width * info->image_height * info->input_components;
  ctx->data = (u8*)malloc(ctx->size);
  if (!ctx->data)
    longjmp(ctx->jump, 1);
  info->dest->next_output_byte = ctx->data;
  info->dest->free_in_buffer = ctx->size;
}

static boolean __JPGWriteAlloc (j_compress_ptr info) noexcept {
  auto ctx = (JpgWriteContext*)info->client_data;
  auto oldSize = ctx->size;
  ctx->size <<= 1;
  ctx->data = (u8*)realloc(ctx->data, ctx->size);
  if (!ctx->data)
    return (boolean)0;
  info->dest->next_output_byte = ctx->data + oldSize;
  info->dest->free_in_buffer = oldSize;
  return (boolean)1;
}

static void __JPGWriteEnd (j_compress_ptr info) noexcept {
  auto ctx = (JpgWriteContext*)info->client_data;
  ctx->iter = ctx->size - info->dest->free_in_buffer;
}

static Texture* __loadJPG (const u8* data, uwl size) noexcept {
  JpgReadContext ctx;
  jpeg_decompress_struct info = {};
  jpeg_error_mgr err;
  u8* pixels = nullptr;
  info.err = jpeg_std_error(&err);
  err.error_exit = __JPGError;

  jpeg_create_decompress(&info);
  info.client_data = &ctx;
  if (setjmp(ctx.jump)) {
    jpeg_destroy_decompress(&info);
    if (pixels)
      free(pixels);
    return nullptr;
  }

  jpeg_mem_src(&info, data, size);
  jpeg_read_header(&info, (boolean)1);

  GLenum fmt;
  if (info.jpeg_color_space == JCS_GRAYSCALE)
    fmt = GL_LUMINANCE;
  else {
    info.out_color_space = JCS_RGB;
    fmt = GL_RGB;
  }
  jpeg_start_decompress(&info);

  auto w = info.output_width;
  auto rowbytes = w * info.output_components;
  auto h = info.output_height;
  auto bytes = rowbytes * h;
  pixels = (u8*)malloc(bytes);
  if (!pixels)
    longjmp(ctx.jump, 1);

  for (auto iter = bytes - rowbytes; info.output_scanline < info.output_height; iter -= rowbytes) {
    JSAMPROW rowPointer = pixels + iter;
    jpeg_read_scanlines(&info, &rowPointer, 1);
  }
  jpeg_finish_decompress(&info);
  jpeg_destroy_decompress(&info);

  auto texture = ObjCache::create<Texture>(w, h, fmt, GL_UNSIGNED_BYTE, pixels);
  free(pixels);
  return texture;
}

static ByteBuffer<> __saveJPG (Texture* texture, s32 quality) noexcept {
  JpgWriteContext ctx = {{}, nullptr, 0, 0};
  jpeg_compress_struct info;
  jpeg_error_mgr err;
  jpeg_destination_mgr dst;
  u8* pixels = nullptr;

  info.err = jpeg_std_error(&err);
  err.error_exit = __JPGError;

  if (setjmp(ctx.jump)) {
    jpeg_destroy_compress(&info);
    if (pixels)
      free(pixels);
    return {};
  }

  auto w = (u32)texture->width();
  auto h = (u32)texture->height();
  auto size = w * h * 3UL;
  pixels = (u8*)malloc(size);
  if (!pixels)
    longjmp(ctx.jump, 1);
  glBindTexture(GL_TEXTURE_2D, texture->id());
  glGetTexImage(GL_TEXTURE_2D, 0, GL_RGB, GL_UNSIGNED_BYTE, pixels);

  jpeg_create_compress(&info);
  info.in_color_space = JCS_RGB;
  info.input_components = 3;
  info.image_width = w;
  info.image_height = h;
  info.client_data = &ctx;
  info.dest = &dst;
  info.dest->init_destination = __JPGWriteStart;
  info.dest->empty_output_buffer = __JPGWriteAlloc;
  info.dest->term_destination = __JPGWriteEnd;

  jpeg_set_defaults(&info);
  jpeg_set_quality(&info, quality, (boolean)1);
  jpeg_start_compress(&info, (boolean)1);

  auto rowBytes = info.image_width * info.input_components;
  auto bytes = rowBytes * info.image_height;
  for (swl iter = bytes - rowBytes; iter >= 0; iter -= rowBytes) {
    JSAMPROW rowPointer = pixels + iter;
    jpeg_write_scanlines(&info, &rowPointer, 1);
  }
  jpeg_finish_compress(&info);
  jpeg_destroy_compress(&info);

  free(pixels);
  return {ctx.data, ctx.size, ctx.iter};
}



Texture* Image::loadImage (ByteBuffer<>& buf) {
  if (!buf.available())
    return nullptr;
  Texture *tex = nullptr;
  auto data = buf.readIter();
  auto size = buf.readable();
  if (isPNG(data, size))
    tex = __loadPNG(data, size);
  else if (isJPG(data, size))
    tex = __loadJPG(data, size);
  return tex;
}

Texture* Image::loadImage (const c8* file) {
  if (!file || !*file)
    return nullptr;
  std::ifstream is;
  is.open(file, std::ios::in | std::ios::binary | std::ios::ate);
  if (is.fail()) {
    log_e("File not exists: %s", file);
    return nullptr;
  }
  auto size = is.tellg();
  if (size < 0)
    return nullptr;
  ByteBuffer<> buf {(uwl)size};
  is.read((c8*)buf.writeIter(), size);
  if (is.fail())
    return nullptr;
  buf.writeIndex((uwl)size);
  return loadImage(buf);
}



ByteBuffer<> Image::saveImage (Texture* texture, bool alpha, int quality) {
  if (!texture || !texture->valid())
    return {};
  return alpha ? __savePNG(texture, quality) : __saveJPG(texture, quality);
}

bool Image::saveImage (Texture* texture, bool alpha, s32 quality, const c8* file) {
  auto buf = saveImage(texture, alpha, quality);
  if (!buf.available())
    return false;
  std::ofstream os;
  os.open(file, std::ios::out | std::ios::binary);
  if (os.fail())
    return false;
  os.write((const c8*)buf.readIter(), buf.readable());
  return os.good();
}


CLOSE_JLIB_NS
